#!/usr/bin/env python
#-*- coding:utf-8 -*-

'''
探测网络主机存活 WEB API
modified by tongtf:

usage:
    http://hostname/ips/ip1|ip2|ip3|ip4|...
result:
    json{ip1:status,ip2:status,ip3:status,...}
'''

import os
import struct
import array
import time
import json
import socket
import threading
from flask import Flask

class SendPingThr(threading.Thread):
    '''
    发送ICMP请求报文的线程
    
    参数：
        ipPool：ip地址池
        icmpPacket：构造icmp报文
        icmpSocket：icmp套接字
        timeout：超时时间
    '''
    
    def __init__(self,ipPool,icmpPacket,icmpSocket,timeout=1):
        threading.Thread.__init__(self)
        self.Sock = icmpSocket
        self.ipPool = ipPool
        self.packet = icmpPacket
        self.timeout = timeout
        self.Sock.settimeout(timeout +1)
        
    def run(self):
        time.sleep(0.01)  #等待接收线程启动
        for ip in self.ipPool:
            try:
                self.Sock.sendto(self.packet,(ip,0))
            except socket.timeout:
                break
        time.sleep(self.timeout)

class Nscan:
    '''
    参数：
        timeout：Socket超时时间，默认1s
        IPv6：默认False
    '''
    
    def __init__(self,timeout=1,IPv6=False):
        self.timeout = timeout
        self.IPv6 = IPv6
        
        self.__data = struct.pack('d',time.time()) #ICMP负荷字节（8bit）
        self.__id = os.getpid() #构造ICMP报文的ID字段，无实际意义
        
    @property
    def __icmpSocket(self):
        '''创建ICMP Socket'''
        if not self.IPv6:
            Sock = socket.socket(socket.AF_INET,socket.SOCK_RAW,socket.getprotobyname('icmp'))
        else:
            Sock = socket.socket(socket.AF_INET,socket.SOCK_RAW,socket.getprotobyname('ipv6-icmp'))
        return Sock
        
    def __inCksum(self,packet):
        '''ICMP报文校验和'''
        if len(packet) & 1:
            packet = packet + '\0'
        words = array.array('h',packet)
        sum = 0
        for word in words:
            sum += (word & 0xffff)
        sum = (sum >> 16) + (sum & 0xffff)
        sum = sum + (sum >> 16)
        
        return (~sum) & 0xffff
        
    @property
    def __icmpPacket(self):
        '''构造ICMP报文'''
        if not self.IPv6:
            header = struct.pack('bbHHh',8,0,0,self.__id,0) #TYPE CODE CHECKSUM ID SEQ
        else:
            header = struct.pack('BbHHh',128,0,0,self.__id,0)
            
        packet = header + self.__data
        chkSum = self.__inCksum(packet)
        
        if not self.IPv6:
            header = struct.pack('bbHHh',8,0,chkSum,self.__id,0)
        else:
            header = struct.pack('BbHHh',128,0,chkSum,self.__id,0)
            
        return header + self.__data
        
    def isUnIP(self,IP):
        '''判断是否为合法单播地址(IPv4)'''
        IP = [int(x) for x in IP.split('.') if x.isdigit()]
        if len(IP) == 4:
            if (0 < IP[0] < 223 and IP[0] != 127 and IP[1] < 256 and IP[2] < 256 and 0 < IP[3] < 255 ):
                return True
        return False

    def chkIpPool(self,ipSet):
        '''过滤ipPool中的非法IP地址'''
        return [ip for ip in ipSet if self.isUnIP(ip)]

    def mPing(self,ipPool):
        '''
        利用ICMP报文探测网络主机存活
        
        参数：
            ipPool：ip地址池
        '''
        Sock = self.__icmpSocket
        Sock.settimeout(self.timeout)
        packet = self.__icmpPacket
        recvFroms = set()
        ip_status = {}
        
        sendThr = SendPingThr(ipPool,packet,Sock,self.timeout)
        sendThr.start()
        
        while True:
            try:
                recvFroms.add(Sock.recvfrom(1024)[1][0])
            except Exception:
                pass
            finally:
                if not sendThr.isAlive():
                    break

        for ip in ipPool:
            if ip in recvFroms:
                ip_status[ip] = 1 #alive
            else:
                ip_status[ip] = 0 #unreachable

        return ip_status

app = Flask(__name__)

'''web入口'''
@app.route('/ips/<ips>')
def ipStatus(ips):
    s = Nscan()
    ipSet = set(unicode(ips).split('|'))
    ipPool = s.chkIpPool(ipSet)
    ip_status = s.mPing(ipPool)
    return json.dumps(ip_status)

'''welcome page'''
@app.route('/')
def index():
    return u'IP Status Check Page'
    
@app.errorhandler(404)
def page_not_found(error):
    return u'access deny'
    
if __name__ == '__main__':
    #app.debug = True
    app.run()